import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import yfinance as yf

# Step 1: 获取历史数据
# 获取三棵树涂料有限公司（603737.SS）过去五年的股价数据
stock_data = yf.download("603737.SS", start="2019-01-01", end="2023-12-31")
close_prices = stock_data['Adj Close']

# 计算每日收益率
daily_returns = close_prices.pct_change().dropna()

# Step 2: 计算年化波动率和平均收益率
mean_return = daily_returns.mean().values[0]  # 提取标量值
volatility = daily_returns.std().values[0]  # 提取标量值

# 将每日收益率年化（假设一年252个交易日）
annual_mean_return = mean_return * 252
annual_volatility = volatility * np.sqrt(252)

# Step 3: 进行蒙特卡洛模拟
# 定义模拟次数和预测时间
simulation_count = 20  # 模拟次数
time_horizon = 252  # 预测未来252个交易日（一年）

# 获取起始价格（历史价格的最后一个值）
start_price = close_prices.iloc[-1]  # 使用 .iloc[-1] 来获取最后一行

# 进行模拟
simulation_matrix = np.zeros((simulation_count, time_horizon))  # 注意矩阵的形状
for i in range(simulation_count):
    # 每次模拟生成一条未来股价路径
    daily_drift = annual_mean_return / 252
    daily_volatility = annual_volatility / np.sqrt(252)
    
    # 随机波动项
    random_shocks = np.random.normal(loc=0, scale=1, size=time_horizon)
    # 生成每日价格变化率
    daily_returns_simulated = np.exp(daily_drift + daily_volatility * random_shocks) - 1
    
    # 计算价格路径
    price_path = np.zeros(time_horizon + 1)
    price_path[0] = start_price
    for j in range(1, time_horizon + 1):
        price_path[j] = price_path[j - 1] * (1 + daily_returns_simulated[j - 1])
    
    # 保存模拟路径
    simulation_matrix[i, :] = price_path[1:]

# Step 4: 计算模拟路径的平均值和边界
mean_simulation = simulation_matrix.mean(axis=0)
upper_bound = np.max(simulation_matrix, axis=0)
lower_bound = np.min(simulation_matrix, axis=0)

# Step 5: 绘制图表
plt.figure(figsize=(14, 8))

# 绘制历史股价曲线
plt.plot(close_prices.index, close_prices, label='Historical Price')

# 创建预测时间序列（未来一年）
future_dates = pd.date_range(start=close_prices.index[-1] + pd.Timedelta(days=1), periods=time_horizon, freq='B')

# 绘制平均模拟路径
plt.plot(future_dates, mean_simulation, color='red', label='Average Simulated Path')

# 绘制上下边界
plt.fill_between(future_dates, lower_bound, upper_bound, color='gray', alpha=0.3, label='Simulation Range')

# 添加图例和标题
plt.legend()
plt.title('603737.SS Stock Price Prediction Using Monte Carlo Simulation')
plt.xlabel('Date')
plt.ylabel('Stock Price')
plt.show()